import json, os, string, ibm_db, requests
import pandas as pd
from typing import Optional

def execute_sql_by_uuid(uuid: str, sql: str, execution_id: Optional[str] = None) -> dict:
    """
    根据UUID执行单条SQL语句
    
    Args:
        uuid: 数据库连接UUID
        sql: 要执行的SQL语句
        execution_id: 执行ID（可选）
        
    Returns:
        dict: 包含执行结果的字典
    """
    conn = None
    try:
        print("*****************************************************************")
        # 获取数据库连接
        conn = TSS_DB_CONNECTION_MANAGER.get_connect(uuid, True)
        
        # 执行SQL语句
        stmt = ibm_db.exec_immediate(conn, sql)
        affected_rows = ibm_db.num_rows(stmt)
        
        # 判断SQL类型
        sql_upper = sql.strip().upper()
        if sql_upper.startswith('SELECT'):
            operation_type = 'QUERY'
            print(f"查询执行完成，执行ID: {execution_id}")
        else:
            operation_type = 'DML'
            print(f"SQL执行完成，影响行数: {affected_rows}, 执行ID: {execution_id}")
        
        # 提交事务（对于DML操作）
        if operation_type == 'DML':
            ibm_db.commit(conn)
            print(f"事务已提交，执行ID: {execution_id}")
        
        return {
            "success": True,
            "operation_type": operation_type,
            "affected_rows": affected_rows,
            "execution_id": execution_id
        }
        
    except Exception as ex:
        # 回滚事务
        if conn:
            ibm_db.rollback(conn)
        print("+" * 50)
        print(f"执行SQL时发生错误: {ex}")
        print("SQL语句:", sql)
        print("+" * 50)
        return {
            "success": False,
            "error": str(ex),
            "execution_id": execution_id,
            "sql": sql
        }
    finally:
        # 关闭数据库连接
        if conn:
            TSS_DB_CONNECTION_MANAGER.release_all_connect()


def insert_dataframe_to_db(uuid: str, table_name: str, df: pd.DataFrame, 
                          truncate_first: bool = True, execution_id: Optional[str] = None) -> dict:
    """
    将DataFrame数据插入到数据库表中
    
    Args:
        uuid: 数据库连接UUID
        table_name: 目标表名（格式：schema.table_name）
        df: 要插入的DataFrame
        truncate_first: 是否先清空表（默认True）
        execution_id: 执行ID（可选）
        
    Returns:
        dict: 包含执行结果的字典
    """
    conn = None
    try:
        print("*****************************************************************")
        # 获取数据库连接
        conn = TSS_DB_CONNECTION_MANAGER.get_connect(uuid, True)
        
        # 第一步：如果需要，先清空表
        if truncate_first:
            truncate_sql = f"DELETE FROM {table_name}"
            stmt = ibm_db.exec_immediate(conn, truncate_sql)
            truncate_count = ibm_db.num_rows(stmt)
            print(f"清空表完成，影响行数: {truncate_count}, 执行ID: {execution_id}")
        
        # 第二步：插入DataFrame数据
        insert_count = 0
        
        # 构建插入SQL - 使用正确的参数占位符
        columns = ', '.join(df.columns)
        # 使用正确的参数占位符（DB2使用:1, :2等）
        placeholders = ', '.join([f":{i+1}" for i in range(len(df.columns))])
        insert_sql = f"INSERT INTO {table_name} ({columns}) VALUES ({placeholders})"
        
        print(f"准备执行插入SQL: {insert_sql}")
        
        # 准备语句
        stmt = ibm_db.prepare(conn, insert_sql)
        if not stmt:
            raise Exception(f"准备SQL语句失败: {ibm_db.stmt_errormsg()}")
        
        # 批量插入数据
        for index, row in df.iterrows():
            # 处理NaN值
            row_data = []
            for col_name, value in row.items():
                if pd.isna(value):
                    row_data.append(None)
                else:
                    # 确保数据类型的兼容性
                    if isinstance(value, (np.int64, np.int32)):
                        value = int(value)
                    elif isinstance(value, (np.float64, np.float32)):
                        value = float(value)
                    row_data.append(value)
            
            # 执行插入 - 使用正确的参数绑定方法
            # 方法1: 直接构建SQL（简单但可能有SQL注入风险，仅用于测试）
            # 方法2: 使用正确的参数绑定
            
            # 这里使用方法2: 构建参数化查询
            try:
                # 使用正确的参数绑定方法
                # 首先尝试使用ibm_db.bind_param（如果可用）
                if hasattr(ibm_db, 'bind_param'):
                    for i, value in enumerate(row_data, 1):
                        ibm_db.bind_param(stmt, i, value)
                    result = ibm_db.execute(stmt)
                else:
                    # 如果bind_param不可用，回退到直接执行（不推荐用于生产）
                    # 构建值字符串
                    values_str = ", ".join([f"'{str(v)}'" if v is not None and not isinstance(v, (int, float)) 
                                           else str(v) if v is not None else "NULL" 
                                           for v in row_data])
                    direct_sql = f"INSERT INTO {table_name} ({columns}) VALUES ({values_str})"
                    result = ibm_db.exec_immediate(conn, direct_sql)
                
                if result:
                    insert_count += 1
                else:
                    error_msg = ibm_db.stmt_errormsg() if hasattr(ibm_db, 'stmt_errormsg') else "未知错误"
                    print(f"插入行失败（索引{index}）: {error_msg}")
                    # 打印失败的行数据前几个字段用于调试
                    print(f"失败行数据（前3个字段）: {dict(list(row.items())[:3])}")
                    
            except Exception as row_ex:
                print(f"插入行时发生错误（索引{index}）: {row_ex}")
                # 继续处理下一行
                continue
            
            # 每插入100行打印一次进度
            if insert_count % 100 == 0:
                print(f"已插入 {insert_count} 行...")
        
        print(f"DataFrame插入完成，插入行数: {insert_count}, 执行ID: {execution_id}")
        
        # 提交事务
        ibm_db.commit(conn)
        print(f"事务已提交，执行ID: {execution_id}")
        
        return {
            "success": True,
            "insert_count": insert_count,
            "execution_id": execution_id
        }
        
    except Exception as ex:
        # 回滚事务
        if conn:
            ibm_db.rollback(conn)
        print("+" * 50)
        print(f"插入DataFrame时发生错误: {ex}")
        print("+" * 50)
        return {
            "success": False,
            "error": str(ex),
            "execution_id": execution_id
        }
    finally:
        # 关闭数据库连接
        if conn:
            TSS_DB_CONNECTION_MANAGER.release_all_connect()

def insert_dataframe_to_db_simple(uuid: str, table_name: str, df: pd.DataFrame, 
                                 truncate_first: bool = True, execution_id: Optional[str] = None) -> dict:
    """
    简化版的DataFrame插入函数，使用批量INSERT语句
    
    Args:
        uuid: 数据库连接UUID
        table_name: 目标表名
        df: 要插入的DataFrame
        truncate_first: 是否先清空表
        execution_id: 执行ID
        
    Returns:
        dict: 执行结果
    """
    conn = None
    try:
        print("*****************************************************************")
        conn = TSS_DB_CONNECTION_MANAGER.get_connect(uuid, True)
        
        # 如果需要，先清空表
        if truncate_first:
            truncate_sql = f"DELETE FROM {table_name}"
            ibm_db.exec_immediate(conn, truncate_sql)
            print(f"清空表完成，执行ID: {execution_id}")
        
        # 构建批量插入的VALUES部分
        columns = ', '.join(df.columns)
        values_list = []
        
        for _, row in df.iterrows():
            row_values = []
            for value in row:
                if pd.isna(value):
                    row_values.append("NULL")
                elif isinstance(value, (int, float)):
                    row_values.append(str(value))
                else:
                    # 转义单引号
                    escaped_value = str(value).replace("'", "''")
                    row_values.append(f"'{escaped_value}'")
            
            values_list.append(f"({', '.join(row_values)})")
        
        # 如果数据量很大，分批插入
        batch_size = 1000
        total_inserted = 0
        
        for i in range(0, len(values_list), batch_size):
            batch_values = values_list[i:i + batch_size]
            values_str = ', '.join(batch_values)
            
            insert_sql = f"INSERT INTO {table_name} ({columns}) VALUES {values_str}"
            
            try:
                stmt = ibm_db.exec_immediate(conn, insert_sql)
                inserted_count = ibm_db.num_rows(stmt)
                total_inserted += inserted_count
                print(f"批量插入进度: {min(i + batch_size, len(values_list))}/{len(values_list)} 行")
            except Exception as batch_ex:
                print(f"批量插入失败（批次 {i//batch_size + 1}）: {batch_ex}")
                # 继续尝试下一批
                continue
        
        # 提交事务
        ibm_db.commit(conn)
        print(f"简化版插入完成，总插入行数: {total_inserted}, 执行ID: {execution_id}")
        
        return {
            "success": True,
            "insert_count": total_inserted,
            "execution_id": execution_id
        }
        
    except Exception as ex:
        if conn:
            ibm_db.rollback(conn)
        print("+" * 50)
        print(f"简化版插入失败: {ex}")
        print("+" * 50)
        return {
            "success": False,
            "error": str(ex),
            "execution_id": execution_id
        }
    finally:
        if conn:
            TSS_DB_CONNECTION_MANAGER.release_all_connect()

# 使用示例
if __name__ == "__main__":
    uuid = "DY_SC_00_LO_ce36a1bfbdbb43858e6cc969ca5e6c0e"
    
    delete_sql = "DELETE FROM BG00MABCDP.T_DWD_GGHT_FJTE_1111"
    result1 = execute_sql_by_uuid(uuid, delete_sql, "delete_execution_001")
    print("删除执行结果:", result1)
    
    # 示例：使用简化版DataFrame插入函数
    sample_data = {
        'SN_NO':["000000000","000000001","000000002"],
        'MedInc': [8.3252, 8.3014, 7.2574],
        'HouseAge': [41.0, 21.0, 52.0],
        'AveRooms': [6.984127, 6.238137, 8.288136],
        'AveBedrms': [1.023810, 0.971880, 1.073446],
        'Population': [322.0, 2401.0, 496.0],
        'AveOccup': [2.555556, 2.109842, 2.802260],
        'Latitude': [37.88, 37.86, 37.85],
        'Longitude': [-122.23, -122.22, -122.24],
        'MedHouseVal': [4.526, 3.585, 3.521]
    }
    sample_df = pd.DataFrame(sample_data)
    
    # 使用简化版函数（推荐）
    result = insert_dataframe_to_db_simple(
        uuid=uuid,
        table_name="BG00MABCDP.T_DWD_GGHT_FJTE_1111",
        df=sample_df,
        truncate_first=True,
        execution_id="simple_dataframe_insert_001"
    )
    print("简化版DataFrame插入结果:", result)
    
    # # 或者尝试使用增强版函数
    # result2 = insert_dataframe_to_db(
    #     uuid=uuid,
    #     table_name="BG00MABCDP.T_DWD_GGHT_FJTE_1111", 
    #     df=sample_df,
    #     truncate_first=False,
    #     execution_id="enhanced_dataframe_insert_001"
    # )
    # print("增强版DataFrame插入结果:", result2)